/*
* Copyright (C) 2010-2012 Magnus Skjegstad
* 
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU General Public License
* as published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
* 
* The latest version of this library can be downloaded from http://mist-pubsub.googlecode.com
* 
*/
package com.skjegstad.mist2.storage;

import com.skjegstad.simpleevents.SimpleAsyncEventListener;
import com.skjegstad.simpleevents.SimpleEvent;
import com.skjegstad.simpleevents.SimpleEventListener;
import com.skjegstad.mist2.MistUtil;
import com.skjegstad.mist2.datatypes.Acknowledgement;
import com.skjegstad.mist2.datatypes.BloomfilterAcknowledgment;
import com.skjegstad.mist2.datatypes.ConnectionIdentifier;
import com.skjegstad.mist2.datatypes.DataElement;
import com.skjegstad.mist2.datatypes.DataFragment;
import com.skjegstad.mist2.datatypes.NodeIdentifier;
import com.skjegstad.mist2.datatypes.SimpleEntry;
import com.skjegstad.mist2.datatypes.TopicIdentifier;
import com.skjegstad.mist2.events.NewDataElementReceivedEvent;
import com.skjegstad.mist2.events.NewDataFragmentsReceived;
import com.skjegstad.mist2.events.RegisterDataElementEvent;
import com.skjegstad.mist2.events.RegisterDataFragmentEvent;
import com.skjegstad.mist2.events.RemoveDataElementsEvent;
import com.skjegstad.mist2.exceptions.MistException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.LinkedBlockingQueue;
import com.skjegstad.mist2.managers.MistEventManager;


/**
 * Event based manager responsible for creating fragments from newly arrived DataElements and make them available
 * to the networking component. As fragments are received, complete fragments are assembled and added to the master data store.
 *
 * @author Magnus Skjegstad
 */
public final class DataFragmentStore {
    private int fragmentSize;
    private final ConcurrentMap<DataElement, List<Long>> completeFragments = new ConcurrentHashMap<DataElement, List<Long>>();
    private final ConcurrentMap<Entry<NodeIdentifier, Integer>, SortedSet<DataFragment>> incompleteFragments = new ConcurrentHashMap();
    private final ConcurrentMap<Long, DataFragment> incompleteHashLookup = new ConcurrentHashMap(); // used to quickly find fragments
    private final ConnectionIdentifier connectionIdentifier;
    private final MistEventManager mistEventManager;
    private final BlockingQueue<Integer> assembleTokenBucket = new LinkedBlockingQueue<Integer>();

    public DataFragmentStore(
            final ConnectionIdentifier connectionIdentifier,
            final int fragmentSize,
            final MistEventManager mistEventManager,
            DataElementStore storeToImport) throws MistException {
        this.fragmentSize = fragmentSize;
        this.mistEventManager = mistEventManager;
        this.connectionIdentifier = connectionIdentifier;

        // import existing elements
        if (storeToImport != null)
            for (DataElement e: storeToImport.getDataElements(storeToImport.getDataElementHashes()))
                this.addDataElementHashes(e);
        
        
        // Listen to events from data manager
        mistEventManager.getDataManagerEventPipe().listen(new SimpleEventListener() {
            public void receiveEvent(SimpleEvent event) {
                // listen to new elements
                if (event instanceof NewDataElementReceivedEvent) {
                    NewDataElementReceivedEvent e = (NewDataElementReceivedEvent) event;
                    //eventManager.triggerEvent(this, new DebugEvent("Adding fragments for element of size " + e.getDataElement().getPayloadLen() + ", frag size " + fragmentSize));
                    addDataElementHashes(e.getDataElement());
                }
                
                // listen to removal of elements
                if (event instanceof RemoveDataElementsEvent) {
                    RemoveDataElementsEvent e = (RemoveDataElementsEvent) event;
                    for (DataElement d : e.getDataElements())
                    completeFragments.remove(d);
                }                
            }
        });

        
        // Listen to events to data fragment manager (us)
        mistEventManager.getDataFragmentEventPipe(this.connectionIdentifier).listen(new SimpleEventListener() {
            public void receiveEvent(SimpleEvent event) {                                
                // listen to arrived fragments (registered)
                if (event instanceof RegisterDataFragmentEvent) {
                    RegisterDataFragmentEvent e = (RegisterDataFragmentEvent) event;

                    if (!e.getConnectionIdentifier().equals(connectionIdentifier)) {
                        MistUtil.debug("Ignored data fragment originating from different connection. Waiting for assembly");
                        return;
                    }

                    List<DataFragment> newFragments = new ArrayList();
                    for (DataFragment f : e.getDataFragments()) {
                        if (addDataFragment(f)) {
                            newFragments.add(f);
                        } else
                            MistUtil.debug("Duplicate or expired fragment received - not added to DataFragmentStore.");
                    }

                    // new fragments are triggered in new event
                    if (!newFragments.isEmpty()) {
                        mistEventManager.getDataFragmentEventPipe(connectionIdentifier).trigger(new NewDataFragmentsReceived(e.getConnectionIdentifier(), newFragments));
                    }
                }
                
                // listen to new data fragments
                if (event instanceof NewDataFragmentsReceived) {
                    try {
                        // attempt to assemble fragments (triggers RegisterDataElementsEvent on complete element, and then NewDataElementReceived on new element)
                        assembleTokenBucket.put(0);
                    } catch (InterruptedException ex) {
                        mistEventManager.throwException(new MistException("Unable to notify fragment assembler", ex));
                        return;
                    }

                }
            }
        });


        // Wait for tokens and assemble fragments async
        mistEventManager.getDataFragmentEventPipe(connectionIdentifier).listen(new SimpleAsyncEventListener() {
            public void receiveEvent(SimpleEvent event) {
                if (event instanceof PeriodicAssembleEvent) {
                    try {
                        assembleTokenBucket.take();
                        while (assembleTokenBucket.size() > 0)
                            assembleTokenBucket.take();
                        assembleFragments();
                        mistEventManager.getDataFragmentEventPipe(connectionIdentifier).trigger(event); // trigger again
                    } catch (Exception ex) {
                        mistEventManager.throwException(new MistException("Got exception while waiting for fragment assembly. Retry in 10 seconds.", ex));
                        try { Thread.sleep(10 * 1000); } catch (InterruptedException ex1) { }
                        mistEventManager.getDataFragmentEventPipe(connectionIdentifier).trigger(new PeriodicAssembleEvent());
                        return;
                    }                                        
                }
            }
        });
        mistEventManager.getDataFragmentEventPipe(connectionIdentifier).trigger(new PeriodicAssembleEvent());
        
    }    

    /**
     * Add hashes from existing data element. The data element is split into fragments and the hash of each fragment is added.
     *
     * @param element
     * @return true on success. False if element already exists.
     */
    protected boolean addDataElementHashes(DataElement element) {
        if (element.isExpired())
            return false;
        
        if (completeFragments.get(element) != null) // check first
            return false;

        // There is a race condition here, as hashes may be generated, inserted and then rejected.
        if (completeFragments.putIfAbsent(element, Collections.synchronizedList(getDataElementHashes(element))) != null)
            return false; // exit if element exist
        
        return true;
    }
    
    /**
     * Get hashes from existing data element. The data element is split into fragments and the hash of each fragment is returned in a list.
     *
     * @param element
     * @return list of hashes for each fragment
     */
    protected LinkedList<Long> getDataElementHashes(DataElement element) {
        if (element.isExpired())
            return new LinkedList();

        LinkedList<Long> hashes = new LinkedList();

        NodeIdentifier sourceNode = element.getSourceNode();
        long seqNr = element.getSeqNr();

        int fragNo = 0;
        for (int start = 0; start < element.getPayloadLen(); start = start + fragmentSize) {
            int len = fragmentSize;
            if (len+start > element.getPayloadLen())
                len = element.getPayloadLen() - start;

            hashes.add(DataElement.getHash(sourceNode, seqNr,
                    fragNo,
                    element.getPayload(),
                    start, len));

            fragNo++;            
        }

        return hashes;
    }
    
    

    /**
     * Add new fragment. The fragment is just added to the receivedFragment list. Call assembleFragments() to
     * look for complete data elements and assemble them.
     *
     * @return true if the fragment is a new fragment, otherwise false.
     */
    protected boolean addDataFragment(DataFragment fragment) {
        if (fragment.isExpired()) {
            MistUtil.debug("Expired data fragment not added to DataFragmentStore (" + fragment.getExpireTime() + ")");
            return false;
        }

        synchronized (incompleteFragments) {
            // get fragments from this source node with the same sequence number
            SortedSet<DataFragment> sortedFragments = incompleteFragments.get(new SimpleEntry(fragment.getSourceNode(), fragment.getSeqNr()));

            // if this is the first fragment, create a new linked list and store it in incompleteFragments map.
            if (sortedFragments == null) {
                sortedFragments = new TreeSet<DataFragment>( new Comparator<DataFragment>() {
                        public int compare(DataFragment o1, DataFragment o2) {
                            if (o1.getFragmentNo() > o2.getFragmentNo())
                                return 1;
                            if (o1.getFragmentNo() < o2.getFragmentNo())
                                return -1;
                            return 0;
                        }
                    });
                incompleteFragments.put(new SimpleEntry(fragment.getSourceNode(), fragment.getSeqNr()), sortedFragments);
            }

            // add the hash to the hashLookup and return true if the hash is new
            if (incompleteHashLookup.putIfAbsent(fragment.getHash(), fragment) == null) {
                // add the fragment to the fragments list
                sortedFragments.add(fragment);
                return true; // new
            }

            MistUtil.debug("Duplicate data fragment not added to DataFragmentStore (" + fragment.getHash() + ")");
            return false; // exists
        }
    }
    
    /**
     * Goes through all received fragments and tries to assemble as many of them as possible. The assembled data elements
     * are removed from the incomplete queue and delivered to the data element store.
     * 
     */
    protected synchronized void assembleFragments() throws MistException {
        List<SortedSet<DataFragment>> candidateFragments = new ArrayList();
        // get lists
        synchronized (incompleteFragments) {
            for (SortedSet<DataFragment> f : incompleteFragments.values())
                if (f != null && 
                    !f.isEmpty() &&
                    f.last().isLastFragment()) {
                        candidateFragments.add(f);
                }
        }

        // sort each list by fragment number
        for (SortedSet<DataFragment> fragmentSet : candidateFragments) {
            synchronized(fragmentSet) {
                // find lists where last is the final fragment and there are no holes
                DataFragment lastFragment = fragmentSet.last();
                
                if (fragmentSet.size() < lastFragment.getFragmentNo()+1) { // number of fragments is larger than length of array, cannot be assembled, continue loop
                    //eventManager.triggerEvent(this, new DebugEvent("Incomplete data element. " + (last.getFragmentNo() - foo.size() + 1) + " elements missing."));

                    /*for (DataFragment f : foo)
                        System.err.println(f.getFragmentNo() + ", " + f.isLastFragment() + ", "+ f.getHash());*/
                    continue;
                }

                if (fragmentSet.size() > lastFragment.getFragmentNo()+1) {
                    throw new MistException("Internal error. Too many fragments to assemble. FragmentSet has " + fragmentSet.size() + " fragments, while last fragment has id " + lastFragment.getFragmentNo());
                }

                // create empty element with blank data of correct size

                DataElement assembled = new DataElement(
                        lastFragment.getSourceNode(),
                        lastFragment.getSeqNr(),
                        new byte[(fragmentSize * (fragmentSet.size()-1)) + lastFragment.getPayloadLen()],
                        lastFragment.getTopic(),
                        lastFragment.getAge(),
                        lastFragment.getExpireTime());

                // fill payload
                int pos = 0;
                int frag = 0;
                boolean failure = false;
                synchronized(fragmentSet) {
                    for (DataFragment newFragment : fragmentSet) {
                        if (newFragment.getFragmentNo() != frag) {// out of order
                            MistUtil.debug("Incomplete data element. Sequence out of order.");
                            failure = true;
                            break;
                        }
                        System.arraycopy(newFragment.getPayload(), 0,
                                assembled.getPayload(), pos, newFragment.getPayloadLen());
                        pos += newFragment.getPayloadLen();
                        frag++;
                    }
                }
                if (failure)
                    continue;

                // if we get here, the assembly was successful.
                MistUtil.debug("Assembled " + fragmentSet.size() + " fragments to one element of " + assembled.getPayloadLen() + " bytes.");

                // remove fragments and hashlookups
                for (DataFragment f : incompleteFragments.get(new SimpleEntry(assembled.getSourceNode(), assembled.getSeqNr())))
                    incompleteHashLookup.remove(f.getHash());
                incompleteFragments.remove(new SimpleEntry(assembled.getSourceNode(), assembled.getSeqNr()));
                
                // trigger event (to data manager)
                mistEventManager.getDataManagerEventPipe().trigger(new RegisterDataElementEvent(assembled, connectionIdentifier));                
            }
        }                
    }

    /**
     * Get hashes of all incomplete data fragments and all complete data elements matching any of the given topics. Complete data elements are added as a single hash to save
     * space.
     *
     * @param topics topics to match.
     * @param excludeExpired do not include expired elements.
     * @param excludeInvalid do not include elements marked as invalid.
     * @return list of hashes w/source nodes matching the topic.
     */
    public Map<Long, NodeIdentifier> getHashesToAck(Collection<TopicIdentifier> topics, boolean excludeExpired, boolean excludeInvalid) {
        HashMap<Long, NodeIdentifier> hashes = new HashMap();

        if (topics == null || topics.isEmpty())
            return hashes; // return empty list if no topics

        // add complete data element fragment hashes first
        synchronized (completeFragments) {
            for (Entry<DataElement, List<Long>> e : completeFragments.entrySet()) {
                if (excludeExpired && e.getKey().isExpired())
                    continue;
                if (excludeInvalid && e.getKey().isInvalid())
                    continue;
                if (e.getKey().getTopic().matchWildcardAny(topics)) {
                    hashes.put(e.getKey().getHash(), e.getKey().getSourceNode());
                    /*for (Long h : e.getValue())
                        hashes.put(h, e.getKey().getSourceNode());*/
                }
            }
        }

        // add hashes of remaining fragments
        synchronized (incompleteHashLookup) {
            for (Entry<Long, DataFragment> e : incompleteHashLookup.entrySet()) {
                if (excludeExpired && e.getValue().isExpired())
                    continue;
                if (excludeInvalid && e.getValue().isInvalid())
                    continue;
                if (e.getValue().getTopic().matchWildcardAny(topics))
                    hashes.put(e.getKey(), e.getValue().getSourceNode());
            }
        }

        return Collections.unmodifiableMap(hashes);
    }

    /**
     * Get list of hashes not in ack, but in topic. Includes complete data elements and fragments from incomplete data elements.
     *
     * @param ack bloom filter acknowledgment.
     * @param topics topics to search in. May include wildcards.
     * @param excludeExpiredDataelements exclude expired data elements
     * @param excludeInvalidDataelements  exclude invalid data elements
     * @return list of hashes (w/source node) not in ack, but in topic.
     */
    public Map<Long, NodeIdentifier> getHashesNotInAck(Acknowledgement ack, 
            Collection<TopicIdentifier> topics, NodeIdentifier excludeNode, 
            boolean excludeExpiredDataelements, boolean excludeInvalidDataelements) {
        Map<Long, NodeIdentifier> hashes = new HashMap();

        if (topics == null || topics.isEmpty())
            return hashes;

        if (ack == null || ack.getLengthInBits() == 0 || ack.getData() == null)
            return this.getHashesToAck(topics, excludeExpiredDataelements, excludeInvalidDataelements); // return everything if no bloomfilter
        
        // Check complete elements first
        for (Entry<DataElement, List<Long>> e : completeFragments.entrySet()) {
            if (excludeExpiredDataelements && e.getKey().isExpired())
                    continue;
            if (excludeInvalidDataelements && e.getKey().isInvalid())
                    continue;
            
            // First, check if the complete data element is acked
            if (!e.getKey().getSourceNode().equals(excludeNode) && !ack.isAcked(e.getKey().getHash())) {
                // Not acked, check all fragments
                for (Long h : e.getValue()) {                    
                    if (!ack.isAcked(h))
                        hashes.put(h, e.getKey().getSourceNode());
                }
            }
        }
        
        // Now check incomplete fragments
        for (Entry<Long, DataFragment> e : incompleteHashLookup.entrySet()) {
            if (excludeExpiredDataelements && e.getValue().isExpired())
                    continue;
            if (excludeInvalidDataelements && e.getValue().isInvalid())
                    continue;
            
            if (!e.getValue().getSourceNode().equals(excludeNode) && 
                    !ack.isAcked(e.getValue().getCompleteHash()) && // complete data element not acked
                    !ack.isAcked(e.getKey())) { // this fragment not acked
                hashes.put(e.getKey(), e.getValue().getSourceNode());
            }
        }
        
        /*
        for (Entry<Long, NodeIdentifier> e : this.getFragmentHashes(topics, excludeExpiredDataelements, excludeInvalidDataelements).entrySet()) {
            if (!e.getValue().equals(excludeNode) && !ack.isAcked(e.getKey()))
                hashes.put(e.getKey(), e.getValue());
        }*/

        return Collections.unmodifiableMap(hashes);
    }    

    /**
     * Get data fragments from a list of hashes. If a hash matches a complete element, all fragments in this element are returned.
     *
     * @param hashes list of hashes.
     * @param removeExpired if expired data elements should be removed from the result.
     * @return data fragments, both complete and incomplete.
     * @throws MistException on error.
     */
    public Collection<DataFragment> getDataFragments(Collection<Long> hashes, boolean removeExpired, boolean excludeInvalid) throws MistException {
        List<DataFragment> ret = new ArrayList();

        for (Long h : hashes) {
            DataFragment d = incompleteHashLookup.get(h);
            if (d == null) {// not in lookup?
                synchronized(completeFragments) { // create fragment from element
                    for (Entry<DataElement, List<Long>> e : completeFragments.entrySet()) {
                        
                        if (removeExpired && e.getKey().isExpired()) // skip expired
                            continue;
                        if (excludeInvalid && e.getKey().isInvalid()) // skip expired
                            continue;
                        
                        int p = e.getValue().indexOf(h); // find correct fragment
                        if (p >= 0) // We found a single fragment with the hash
                            ret.add(new DataFragment(e.getKey(), fragmentSize, p));           
                        
                        // Check if hash matches complete element
/*                        if (e.getKey().getHash() == h) {
                            // Add all fragments from element
                            for (int i = 0; i < e.getValue().size(); i++) {
                                ret.add(new DataFragment(e.getKey(), fragmentSize, i));
                            }
                        }*/
                    }
                }
            } else {
                if (removeExpired && d.isExpired())
                    continue;
                if (excludeInvalid && d.isInvalid())
                    continue;
                ret.add(d); // add incomplete fragment
            }
        }

        return ret;
    }
    
    private class PeriodicAssembleEvent implements SimpleEvent {}
}
